{
  "cells": [
    {
      "metadata": {
        "id": "yR8zLEl-Q-sR"
      },
      "cell_type": "markdown",
      "source": [
        "This notebook is a tutorial that demonstrates how to configure a questionnaire simulation using Concordia."
      ]
    },
    {
      "metadata": {
        "id": "MczTJA7NQ9VM"
      },
      "cell_type": "markdown",
      "source": [
        "# <a href=\"https://colab.research.google.com/github/google-deepmind/concordia/blob/main/examples/questionnaire_example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "DBOYi3m8SBPa"
      },
      "cell_type": "code",
      "source": [
        "# @title Colab-specific setup (use a CodeSpace to avoid the need for this).\n",
        "try:\n",
        "  %env COLAB_RELEASE_TAG\n",
        "except:\n",
        "  pass  # Not running in colab.\n",
        "else:\n",
        "  %pip install --ignore-requires-python --requirement 'https://raw.githubusercontent.com/google-deepmind/concordia/main/examples/requirements.in' 'git+https://github.com/google-deepmind/concordia.git#egg=gdm-concordia'\n",
        "  %pip list"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "psW9YfYltDzj"
      },
      "cell_type": "code",
      "source": [
        "# @title Imports\n",
        "\n",
        "from concordia.contrib import language_models as language_model_utils\n",
        "import concordia.prefabs.entity as entity_prefabs\n",
        "import concordia.prefabs.game_master as game_master_prefabs\n",
        "import concordia.prefabs.simulation.questionnaire_simulation as questionnaire_simulation\n",
        "from concordia.typing import prefab as prefab_lib\n",
        "from concordia.utils import helper_functions\n",
        "from IPython import display\n",
        "import numpy as np\n",
        "import sentence_transformers"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "Nin9rPBbcM0R"
      },
      "cell_type": "code",
      "source": [
        "# @title Language Model Selection: provide key or select DISABLE_LANGUAGE_MODEL\n",
        "\n",
        "# By default this colab uses models via an external API so you must provide an\n",
        "# API key. TogetherAI offers open weights models from all sources.\n",
        "\n",
        "API_KEY = ''  # @param {type: 'string'}\n",
        "# See concordia/language_model/utils.py\n",
        "API_TYPE = 'openai'  # e.g. 'together_ai' or 'openai'.\n",
        "MODEL_NAME = (  # for API_TYPE = 'together_ai', we recommend MODEL_NAME = 'google/gemma-3-27b-it'\n",
        "    'gpt-5'\n",
        ")\n",
        "# To debug without spending money on API calls, set DISABLE_LANGUAGE_MODEL=True\n",
        "DISABLE_LANGUAGE_MODEL = False"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "vnxj8_ce5eFd"
      },
      "cell_type": "code",
      "source": [
        "# @title Use the selected language model\n",
        "\n",
        "# Note that it is also possible to use local models or other API models,\n",
        "# simply replace this cell with the correct initialization for the model\n",
        "# you want to use.\n",
        "\n",
        "if not DISABLE_LANGUAGE_MODEL and not API_KEY:\n",
        "  raise ValueError('API_KEY is required.')\n",
        "\n",
        "model = language_model_utils.language_model_setup(\n",
        "    api_type=API_TYPE,\n",
        "    model_name=MODEL_NAME,\n",
        "    api_key=API_KEY,\n",
        "    disable_language_model=DISABLE_LANGUAGE_MODEL,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "brdgSD2NuwOQ"
      },
      "cell_type": "markdown",
      "source": [
        "## Language Model setup"
      ]
    },
    {
      "metadata": {
        "id": "Ez6153pSuwOQ"
      },
      "cell_type": "code",
      "source": [
        "# @title Setup sentence encoder\n",
        "\n",
        "if DISABLE_LANGUAGE_MODEL:\n",
        "  embedder = np.ones(3)\n",
        "else:\n",
        "  st_model = sentence_transformers.SentenceTransformer(\n",
        "      'sentence-transformers/all-mpnet-base-v2'\n",
        "  )\n",
        "  embedder = lambda x: st_model.encode(x, show_progress_bar=False)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "_j3qaLBNHp_P"
      },
      "cell_type": "code",
      "source": [
        "test = model.sample_text(\n",
        "    'Is societal and technological progress like getting a clearer picture of '\n",
        "    'something true and deep?'\n",
        ")\n",
        "print(test)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "ixIOLiWL4s8-"
      },
      "cell_type": "markdown",
      "source": [
        "# Parameters"
      ]
    },
    {
      "metadata": {
        "id": "2jg5UkxprRCz"
      },
      "cell_type": "code",
      "source": [
        "# @title Load prefabs from packages to make the specific palette to use here.\n",
        "\n",
        "prefabs = {\n",
        "    **helper_functions.get_package_classes(entity_prefabs),\n",
        "    **helper_functions.get_package_classes(game_master_prefabs),\n",
        "}"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "7yK5YJqsJe-R"
      },
      "cell_type": "code",
      "source": [
        "display.display(\n",
        "    display.Markdown(helper_functions.print_pretty_prefabs(prefabs))\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "l77tYjIl3xBo"
      },
      "cell_type": "code",
      "source": [
        "from concordia.contrib.data.questionnaires import depression_anxiety_stress_scale\n",
        "\n",
        "DASS = depression_anxiety_stress_scale.DASSQuestionnaire()"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "1ojL80JBdfcy"
      },
      "cell_type": "code",
      "source": [
        "lord_of_the_rings_characters = [\n",
        "    \"Frodo Baggins\",\n",
        "    \"Samwise Gamgee\",\n",
        "    \"Meriadoc Brandybuck\",\n",
        "    \"Peregrin Took\",\n",
        "    \"Gandalf\",\n",
        "    \"Aragorn\",\n",
        "    \"Legolas\",\n",
        "    \"Gimli\",\n",
        "    \"Boromir\",\n",
        "    \"Sauron\",\n",
        "    \"Saruman\",\n",
        "    \"Gollum\",\n",
        "    \"Arwen\",\n",
        "    \"Galadriel\",\n",
        "    \"Elrond\",\n",
        "    \"Théoden\",\n",
        "    \"Éowyn\",\n",
        "    \"Faramir\",\n",
        "    \"Denethor\",\n",
        "    \"Bilbo Baggins\",\n",
        "]"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "0g3grivzdhJm"
      },
      "cell_type": "code",
      "source": [
        "# @title Scene 1 - Frodo in the Shire\n",
        "\n",
        "instances = []\n",
        "\n",
        "for name in lord_of_the_rings_characters:\n",
        "  instances.append(\n",
        "      prefab_lib.InstanceConfig(\n",
        "          prefab='basic__Entity',\n",
        "          role=prefab_lib.Role.ENTITY,\n",
        "          params={\n",
        "              'name': name,\n",
        "          },\n",
        "      )\n",
        "  )\n",
        "\n",
        "instances.append(\n",
        "    prefab_lib.InstanceConfig(\n",
        "        prefab='interviewer__GameMaster',\n",
        "        role=prefab_lib.Role.GAME_MASTER,\n",
        "        params={\n",
        "            'name': 'interviewer',\n",
        "            'player_names': lord_of_the_rings_characters,\n",
        "            'questionnaires': [DASS],\n",
        "            'verbose': False,\n",
        "            'embedder': embedder,\n",
        "        },\n",
        "    ),\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "68eqdgXZMneJ"
      },
      "cell_type": "code",
      "source": [
        "config = prefab_lib.Config(\n",
        "    default_premise=(\n",
        "        f'Frodo is making his way to mount Doom, while Sauron is searching for'\n",
        "        f' the Ring.'\n",
        "    ),\n",
        "    default_max_steps=1,\n",
        "    prefabs=prefabs,\n",
        "    instances=instances,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "GqzRE7oOYN_m"
      },
      "cell_type": "code",
      "source": [
        "# @title Initialize the simulation\n",
        "runnable_simulation = questionnaire_simulation.QuestionnaireSimulation(\n",
        "    config=config,\n",
        "    model=model,\n",
        "    embedder=embedder,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "M4Z1ttTfuwOR"
      },
      "cell_type": "code",
      "source": [
        "# @title Run the simulation\n",
        "results_log = runnable_simulation.play(max_steps=1)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "pewMIIPmNYdT"
      },
      "cell_type": "code",
      "source": [
        "questionnaire_comp = runnable_simulation.game_masters[0].get_component(\n",
        "    'questionnaire'\n",
        ")\n",
        "questionnaire_result_df = questionnaire_comp.get_questionnaires_results()\n",
        "print(questionnaire_result_df)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "ynq0s6iQj_ew"
      },
      "cell_type": "code",
      "source": [
        "print(questionnaire_comp.get_answers())"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "8DovnYWTIl-Y"
      },
      "cell_type": "markdown",
      "source": [
        "```\n",
        "Copyright 2025 DeepMind Technologies Limited.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python:ml_python_notebook",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
